import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from torch.autograd import Variable

_EPS = 1e-10


class MLP(nn.Module):
    """Two-layer fully-connected ELU net with batch norm."""

    def __init__(self, n_in, n_hid, n_out, do_prob=0.):
        super(MLP, self).__init__()
        self.fc1 = nn.Linear(n_in, n_hid)
        self.fc2 = nn.Linear(n_hid, n_out)
        self.bn = nn.BatchNorm1d(n_out)
        self.dropout_prob = do_prob

        self.init_weights()

    def init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.xavier_normal_(m.weight.data)
                m.bias.data.fill_(0.1)
            elif isinstance(m, nn.BatchNorm1d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()


    def forward(self, inputs):
        # Input shape: [num_sims, num_things, num_features]
        x = F.elu(self.fc1(inputs))
        x = F.dropout(x, self.dropout_prob, training=self.training)
        x = F.elu(self.fc2(x))
        return self.bn(x)



class LossFunc(nn.Module):
    def __init__(self, weights):
        super(LossFunc, self).__init__()
        self.weights = weights

    def forward(self, predictions, targets):
        # 计算加权均方误差
        loss = torch.abs(predictions - targets).mean(dim=0).mean(dim=1) 
        loss = loss * self.weights
        loss = torch.mean(loss)
        return loss

def WMAPE(target, predict):
    WMAPEs = []
    for i in range(target.shape[0]):
        wmapei = torch.abs(target[i] - predict[i])/target[i] 
        weights = target[i]/torch.sum(target[i],dim=0)
        wmapei = wmapei * weights
        wmapei = torch.sum(wmapei,dim=0)
        WMAPEs.append(wmapei.view(1,2))
    WMAPEs = torch.cat(WMAPEs, dim=0)
    weights = target.sum(dim=1)
    weights = weights/weights.sum(dim=0)
    WMAPE = torch.mean(WMAPEs*weights, dim=0)
    return WMAPE


class UsagiNet(nn.Module):
    def __init__(self, input_size, hidden_size=512, output_size=14, do_prob=0.):
        super(UsagiNet, self).__init__()
        self.mlp0 = MLP(hidden_size, int(hidden_size), hidden_size, do_prob)
        self.mlp1 = MLP(hidden_size, hidden_size, output_size, do_prob)
        self.gru = nn.GRU(input_size, hidden_size, 1, batch_first=True)
        
    def forward(self, x1, x2):
        h1, _ = self.gru(x1)
        h2  = self.mlp0(x2)
        h = F.relu(h1[:,-1,:]+h2)
        out = self.mlp1(h) 
        return out

'''
class UsagiNet(nn.Module):
    def __init__(self, input_size, hidden_size=512, output_size=14, do_prob=0.):
        super(UsagiNet, self).__init__()
        self.mlp0 = MLP(input_size, hidden_size, hidden_size, do_prob)
        self.mlp1 = MLP(hidden_size, hidden_size, output_size, do_prob)
        self.gru = nn.GRU(hidden_size, hidden_size, 1, batch_first=True)
        

    def forward(self, x):
        size = x.shape
        x = x.reshape(size[0]*size[1],-1)
        out = self.mlp0(x)
        out = out.reshape(size[0],size[1],-1).contiguous()
        out, _ = self.gru(out)
        out = self.mlp1(out[:, -1, :])  # 取GRU最后一时间步的输出
        return out
''' 
    